#!user/bin/env python3
# -- coding: utf-8 --
# @Createtime：2022/5/31
# @Updatetime：11:56
# @Author：Test008
# @File：baidu_ocr.py
# @Description：百度OCR图片文字识别工具

# encoding:utf-8
import requests
import os


def gettoken():
    # client_id 为官网获取的AK， client_secret 为官网获取的SK
    client_id = '0NQdDGdba81iX5A5GGpY5oPU'
    client_secret = 'wPVA3TzjfPq765z1ISnhGAI8jYxi7Xcp'
    host = f'https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={client_id}&client_secret={client_secret}'
    response = requests.get(host)
    if response:
        print(response.json())
        return response.json()['access_token']


def getnum(imagedata: str = None):
    import requests
    import base64

    '''
    通用文字识别
    '''

    request_url = "https://aip.baidubce.com/rest/2.0/ocr/v1/general_basic"
    # 二进制方式打开图片文件
    if imagedata:
        img = imagedata.encode()
    else:
        import os
        path = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'datafiles')
        f = open(path+'/num.png', 'rb')
        img = base64.b64encode(f.read())

    params = {"image": img}
    access_token = gettoken()
    request_url = request_url + "?access_token=" + access_token
    headers = {'content-type': 'application/x-www-form-urlencoded'}
    response = requests.post(request_url, data=params, headers=headers)
    if response:
        print(response.json())
        # 异常处理
        try:
            return response.json()["words_result"][0]['words']
        except IndexError:
            return "验证码获取失败"


if __name__ == '__main__':
    getnum()
